import numpy as np
from scipy.fftpack import next_fast_len
from scipy.fft import fft, ifft, fftfreq
import scipy.signal as signal
import os

def z_score_normalization(data):
    mean = np.mean(data, axis=0, keepdims=True)
    std = np.std(data, axis=0, keepdims=True)
    return (data - mean) / std

# 定义Hjorth参数的计算函数，计算脑电信号的活动性、移动性和复杂性
def calculate_hjorth_parameters(segment):
    mobility = np.std(np.diff(segment)) / np.std(segment)
    complexity = np.std(np.diff(np.diff(segment))) / np.std(np.diff(segment))
    return np.var(segment), mobility, complexity / mobility


# 谱熵计算函数
def spectral_entropy(fft_result):
    # 计算功率谱
    power_spectrum = np.abs(fft_result)**2
    # 计算总功率
    total_power = np.sum(power_spectrum)
    # 计算概率分布
    probability_distribution = power_spectrum / total_power
    eps = np.finfo(float).eps
    # 计算 P(m) log2 P(m) 的加和
    return -np.sum(probability_distribution * np.log2(probability_distribution + eps))

hilbert3 = lambda x: signal.hilbert(x, next_fast_len(len(x)),axis=0)[:len(x)]



if __name__ == '__main__':
    # 划窗，窗口长度为1s
    winL = 15
    duration = 15
    fs = 400
    num_segments = duration // winL
    root_data_path = '/root/data/video_decoding'
    eeg_path = os.path.join(root_data_path, 'cut_data_zscore')      # original file path
    outputs = os.path.join(root_data_path, 'npy_data_class_2')   # output feature
    time_files = os.listdir(eeg_path)
    
    # 定义波段
    # TODO: high gamma > 120 Hz 信号可能无效，high_gamma 占比过高可能是假信号
    max_band = max(fs, 120)
    bands = {
            'theta': (4, 7),
            'alpha': (8, 12),
            'beta': (13, 30),
            'low_gamma': (31, 70),
            'high_gamma': (70, max_band)
        }
    for data_id in ['train', 'val']:
        file_names = os.listdir(os.path.join(eeg_path, data_id))
        dir_path = f'{outputs}/{data_id}'
        for file in file_names:
            print(file)
            if file[16]=='1':
                class_start = 0
            elif file[16]=='2':
                class_start = 2
            # neural data
            seeg_data = np.load(os.path.join(eeg_path, data_id, file))  # read data
            seeg_data = seeg_data.transpose((1, 0))

            # 增加数据增强步骤：随机添加小幅度噪声
            noise_level = 0.1  # 可调整的噪声水平
            noisy_data = seeg_data + np.random.normal(0, noise_level, seeg_data.shape)
            # 去除线性趋势
            segments = signal.detrend(noisy_data, axis=0)
            segments = np.array_split(segments, num_segments)

            for segment in segments:
                # 频率分辨率
                freqs = fftfreq(len(segment), d=1/fs)
                # 快速傅里叶变换
                fft_result = fft(segment)
                amplitude_spectrum = np.abs(fft_result)

                band_features = []
                for band in bands:
                    feature = []
                    low, high = bands[band]
                    # 计算能量值
                    idx = np.where((freqs >= low) & (freqs <= high))[0]
                    sig_energy = np.sum(amplitude_spectrum[idx]**2) / len(idx)

                    band_filter = (freqs >= low) & (freqs <= high)
                    filtered_fft_values = np.zeros_like(fft_result)
                    filtered_fft_values[band_filter] = fft_result[band_filter]
                    
                    # 谱熵在FFT后的信号上计算
                    spec_entropy = spectral_entropy(filtered_fft_values)  

                    # 逆傅里叶变换 (time domain)
                    filtered_signal = np.real(ifft(filtered_fft_values))
                    
                    # Hjorth参数在重构信号上计算
                    var, mobi, comp = calculate_hjorth_parameters(filtered_signal)
                    
                    feature.append(var)
                    feature.append(mobi)
                    feature.append(comp)
                    feature.append(spec_entropy)
                    feature.append(sig_energy)
                    band_features.append(feature)
                
                np.save(os.path.join(dir_path, file[:-4]+'_'+str(class_start)), band_features)
                class_start += 1